// SPDX-License-Identifier: GPL-2.0

#define pr_fmt(fmt) KBUILD_MODNAME ": " fmt

#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/fs.h>
#include <linux/err.h>
#include <linux/errno.h>
#include <linux/slab.h>
#include <linux/uaccess.h>
#include <linux/limits.h>
#include <linux/sched/signal.h>
#include <linux/utsname.h>
#include <linux/highmem.h>
#include <linux/vmalloc.h>
#include <linux/sched/mm.h>

#include "dim_core_dynamic_baseline.h"
#include "dim_core_static_baseline.h"
#include "dim_core_mem.h"
#include "dim_core_status.h"
#include "dim_core_policy.h"
#include "dim_core_measure_log.h"
#include "../common/dim_calculate.h"
#include "../common/dim_utils.h"
#include "../common/dim_pages.h"
#include "../adapter/dim_adapter.h"

/* get kernel module list */
extern struct list_head *get_kernel_modules(void);

static struct dim_baseline_handle *dim_dynamic_baseline;

static int dim_dynamic_baseline_compare(void *data1, void *data2)
{
	return dim_baseline_node_compare((struct dim_baseline_node *)data1,
					 (struct dim_baseline_node *)data2);
}

static void digest_extend_and_write(char *name,
				    char *digest,
				    enum dim_hash_type hash_index)
{
	int ret;

	ret = dim_extend(digest, dim_hash_get_digest_len(hash_index));
	if (ret < 0) {
		pr_err("extend pcr failed, %s, %d\n", name, ret);
		return;
	}

	ret = dim_write_measure_log(name, digest, hash_index);
	if (ret != 0)
		pr_err("write to measure log failed, %s, %d\n", name, ret);
}

void dim_measure_baseline_record(char *name,
				 const char *err_info,
				 char *digest,
				 enum dim_hash_type hash_index,
				 int search_all)
{
	char *name_buf = NULL;
	char *digest_buf = digest;

	if (name == NULL || hash_index >= DIM_HASH_LAST ||
	    (err_info == NULL && digest == NULL))
		return;

	if (err_info == NULL) {
		name_buf = kstrdup(name, GFP_KERNEL);
		if (name_buf == NULL)
			return;
	} else {
		name_buf = kmalloc(strlen(name) + strlen(err_info) + 3,
			GFP_KERNEL);
		if (name_buf == NULL)
			return;
		sprintf(name_buf, "%s, %s", name, err_info);
	}

	if (digest_buf == NULL) {
		digest_buf = kzalloc(MAX_DIGEST_SIZE, GFP_KERNEL);
		if (digest_buf == NULL) {
			kfree(name_buf);
			return;
		}
	}

	if (!dim_measure_log_find_baseline(name_buf, digest_buf,
		hash_index, search_all))
		digest_extend_and_write(name_buf, digest_buf, hash_index);

	kfree(name_buf);
	if (digest_buf != digest)
		kfree(digest_buf);
}

static int
dim_dynamic_baseline_create_and_add_node(char *name,
					 enum dim_baseline_type type,
					 enum dim_hash_type algo,
					 unsigned char *data,
					 unsigned int data_len)
{
	struct dim_baseline_node *node = NULL;
	int ret;
	enum dim_core_policy_action action = ACTION_LOG;

	if (name == NULL || data == NULL)
		return -EINVAL;

	node = dim_baseline_node_create(dim_dynamic_baseline);
	if (IS_ERR(node))
		return PTR_ERR(node);

	action = dim_policy_get_action(name);

	do {
		ret = dim_baseline_node_set_filename(dim_dynamic_baseline, node, name);
		if (ret)
			break;

		ret = dim_baseline_node_set_type(dim_dynamic_baseline, node, type);
		if (ret)
			break;

		ret = dim_baseline_node_set_hash_algo(dim_dynamic_baseline, node, algo);
		if (ret)
			break;

		ret = dim_baseline_node_set_action(dim_dynamic_baseline, node, action);
		if (ret)
			break;

		ret = dim_baseline_node_set_main_hash_data(dim_dynamic_baseline, node,
						data, data_len);
		if (ret)
			break;
	} while (0);

	if (ret) {
		dim_baseline_node_destroy(dim_dynamic_baseline, node);
		return ret;
	}

	ret = dim_baseline_add_node(dim_dynamic_baseline, node);
	if (ret < 0) {
		dim_baseline_node_destroy(dim_dynamic_baseline, node);
		return ret;
	}

	return 0;
}

static int
dim_dynamic_baseline_user_create_trampoline(struct mm_struct *mm,
					    struct vm_area_struct *vma,
					    char *path)
{
	int ret, hash_len;
	struct dim_baseline_node *find_node = NULL;
	size_t name_len = strlen(path) + strlen(DIM_TRAMPOLINE_NAME_SUFFIX);
	char *name_buf = NULL;
	char digest_out[MAX_DIGEST_SIZE];

	if (name_len > 65536U)
		return -ENAMETOOLONG;

	/* get static baseline data and record */
	find_node = dim_static_baseline_search(path, DIM_BASELINE_TRAMPOLINE);
	if (find_node == NULL)
		return 0; /* no static baseline, needn't create */

	name_buf = kmalloc(name_len + 1, GFP_KERNEL);
	if (name_buf == NULL)
		return -ENOMEM;
	snprintf(name_buf, name_len + 1, "%s%s",
		path, DIM_TRAMPOLINE_NAME_SUFFIX);

	dim_measure_baseline_record(name_buf, NULL,
		find_node->data, find_node->algo, 1);

	/* calculate and compare digest */
	hash_len = dim_hash_get_digest_len(find_node->algo);
	ret = dim_calc_user_task_trampoline_digest(mm, vma,
		find_node->algo, digest_out);
	if (ret != hash_len) {
		pr_warn("fail to calculate trampoline of %s: %d\n",
			path, ret);
	} else if (dim_digest_compare(find_node->data,
		digest_out, hash_len) != 0) {
		dim_measure_baseline_record(name_buf, "tampered!",
			digest_out, find_node->algo, 0);
	}

	/* create dynamic baseline */
	if (find_node->algo != dim_hash_get_cur_index()) {
		hash_len = dim_hash_get_digest_len(dim_hash_get_cur_index());
		ret = dim_calc_user_task_trampoline_digest(mm, vma,
			dim_hash_get_cur_index(), digest_out);
		if (ret != hash_len) { /* calculate failed */
			pr_warn("fail to calculate trampoline of %s: %d\n",
				path, ret);
			kfree(name_buf);
			return ret < 0 ? ret : -EINVAL;
		}
	}

	ret = dim_dynamic_baseline_create_and_add_node(path,
		DIM_BASELINE_TRAMPOLINE, dim_hash_get_cur_index(),
		digest_out, hash_len);
	if (ret < 0)
		pr_err("fail to add trampoline baseline of %s: %d\n",
			path, ret);

	kfree(name_buf);
	return ret;
}

static int dim_dynamic_baseline_user_create_text(struct mm_struct *mm,
						 struct vm_area_struct *vma,
						 char *path)
{
	int ret, hash_len;
	struct dim_baseline_node *find_node = NULL;
	char digest_out[MAX_DIGEST_SIZE];

	/* get static baseline data and record */
	find_node = dim_static_baseline_search(path, DIM_BASELINE_USER);
	if (find_node == NULL) {
		dim_measure_baseline_record(path, "no static baseline",
			NULL, dim_hash_get_cur_index(), 1);
	} else {
		dim_measure_baseline_record(path, NULL,
			find_node->data, find_node->algo, 1);
		/* calculate and compare digest */
		hash_len = dim_hash_get_digest_len(find_node->algo);
		ret = dim_calc_user_task_text_digest(mm, vma,
			find_node->algo, digest_out);
		if (ret != hash_len) {
			pr_warn("fail to calculate digest of %s: %d\n",
				path, ret);
			return ret < 0 ? ret : -EINVAL;
		}

		if (dim_digest_compare(find_node->data,
			digest_out, hash_len) != 0)
			dim_measure_baseline_record(path, "tampered!",
				digest_out, find_node->algo, 0);
	}

	/* create dynamic baseline */
	if (find_node == NULL ||
	    find_node->algo != dim_hash_get_cur_index()) {
		hash_len = dim_hash_get_digest_len(dim_hash_get_cur_index());
		ret = dim_calc_user_task_text_digest(mm, vma,
			dim_hash_get_cur_index(), digest_out);
		if (ret != hash_len) { /* calculate failed */
			pr_warn("fail to calculate digest of %s: %d\n",
				path, ret);
			return ret < 0 ? ret : -EINVAL;
		}
	}

	ret = dim_dynamic_baseline_create_and_add_node(path,
		DIM_BASELINE_USER, dim_hash_get_cur_index(),
		digest_out, hash_len);
	if (ret < 0)
		pr_err("fail to add baseline of %s: %d\n", path, ret);

	return ret;
}

static int
dim_dynamic_baseline_user_create_from_vma(struct mm_struct *mm,
					  struct vm_area_struct *vma,
					  char *path)
{
	int ret;

	if (!dim_policy_has_item(path) ||
	    dim_dynamic_baseline_search(path, DIM_BASELINE_USER) != NULL)
		return 0;

	ret = dim_dynamic_baseline_user_create_trampoline(mm, vma, path);
	if (ret < 0) {
		pr_warn("fail to create tramp dynamic baseline of %s: %d\n",
			path, ret);
		return ret == -EINTR ? ret : 0;
	}

	ret = dim_dynamic_baseline_user_create_text(mm, vma, path);
	if (ret < 0) {
		pr_warn("fail to create dynamic baseline of %s: %d\n",
			path, ret);
		return ret == -EINTR ? ret : 0;
	}

	return 0;
}

struct measure_tsk {
	struct task_struct *t;
	struct list_head list;
};

static void add_to_measure_list(struct list_head *mtsks,
				struct task_struct *tsk)
{
	struct measure_tsk *mtsk = NULL;

	mtsk = kmalloc(sizeof(struct measure_tsk), GFP_ATOMIC);
	if (mtsk == NULL)
		return;

	mtsk->t = tsk;
	list_add_tail(&mtsk->list, mtsks);
}

static int dim_dynamic_baseline_user_create(void)
{
	struct task_struct *tsk = NULL;
	struct vm_area_struct *vma = NULL;
	struct mm_struct *mm = NULL;
	char *path = NULL;
	char *path_buf = NULL;
	struct file *vm_file_prev = NULL;
	struct list_head mtsks;
	struct measure_tsk *pos = NULL;
	struct measure_tsk *next = NULL;
	int ret;

	INIT_LIST_HEAD(&mtsks);

	path_buf = kmalloc(PATH_MAX, GFP_KERNEL);
	if (path_buf == NULL)
		return -ENOMEM;

	/* get all tasks, add every task's usage and add to measure task list */
	read_lock(&tasklist_lock);
	for_each_process(tsk) {
		get_task_struct(tsk);
		add_to_measure_list(&mtsks, tsk);
	}
	read_unlock(&tasklist_lock);

	/* do measure from measure task list */
	list_for_each_entry_safe(pos, next, &mtsks, list) {
		tsk = pos->t;
		mm = get_task_mm(tsk);
		if (mm == NULL) {
			put_task_struct(tsk);
			kfree(pos);
			continue;
		}

		down_read(&mm->mmap_lock);
		for (vma = mm->mmap; vma != NULL; vma = vma->vm_next) {
			if (!(vma->vm_flags & VM_READ) ||
			    !(vma->vm_flags & VM_EXEC) ||
			    (vma->vm_flags & VM_WRITE) ||
			    dim_get_vma_file(vma) == NULL ||
			    dim_get_vma_file(vma) == vm_file_prev)
				continue;

			vm_file_prev = dim_get_vma_file(vma);

			/* get elf file path from vma */
			path = d_path(&vm_file_prev->f_path,
				path_buf, PATH_MAX);
			if (IS_ERR(path)) {
				pr_warn("fail to get vm file path: %ld\n",
					PTR_ERR(path));
				continue;
			}

			ret = dim_dynamic_baseline_user_create_from_vma(mm,
				vma, path);
			if (ret < 0)
				break;
		}
		up_read(&mm->mmap_lock);
		mmput(mm);
		put_task_struct(tsk);
		kfree(pos);
	}

	kfree(path_buf);
	return 0;
}

static int dim_dynamic_baseline_ko_create(void)
{
	struct module *mod = NULL;
	struct list_head *head = get_kernel_modules();
	char digest_out[MAX_DIGEST_SIZE];
	int ret, hash_len;
	char ko_name[NAME_MAX * 2 + 1];
	struct dim_baseline_node *find_node = NULL;

	if (head == NULL)
		return -EFAULT;

	mutex_lock(&module_mutex);
	list_for_each_entry_rcu(mod, head, list) {
		if (mod->state == MODULE_STATE_UNFORMED ||
		    mod->state != MODULE_STATE_LIVE ||
		    !dim_policy_has_item(mod->name))
			continue;

		/* search if the dynamic baseline has been created */
		if (strlen(init_uts_ns.name.release) + strlen(mod->name) + 1 >= sizeof(ko_name)) {
			pr_warn("overlength of %s or its kernel release\n",
				ko_name);
			continue;
		}
		sprintf(ko_name, "%s/%s",
			init_uts_ns.name.release, mod->name);

		if (dim_dynamic_baseline_search(ko_name,
			DIM_BASELINE_KERNEL) != NULL)
			continue;

		find_node = dim_static_baseline_search(ko_name,
			DIM_BASELINE_KERNEL);
		if (find_node == NULL) {
			pr_warn("can't find static baseline of %s\n",
				ko_name);
			dim_measure_baseline_record(ko_name,
				"no static baseline",
				NULL, dim_hash_get_cur_index(), 1);
		} else {
			/* output to measure log and pcr */
			dim_measure_baseline_record(ko_name, NULL,
				find_node->data, find_node->algo, 1);
		}

		/* create and add dynamic baseline */
		hash_len = dim_hash_get_digest_len(dim_hash_get_cur_index());
		ret = dim_calc_ko_digest(mod,
			dim_hash_get_cur_index(), digest_out);
		if (ret != hash_len) {
			pr_warn("fail to calculate %s digest: %d\n",
				ko_name, ret);
			continue;
		}

		ret = dim_dynamic_baseline_create_and_add_node(ko_name,
			DIM_BASELINE_KERNEL, dim_hash_get_cur_index(),
			digest_out, hash_len);
		if (ret != 0) {
			pr_err("fail to add baseline of %s: %d\n",
				ko_name, ret);
			continue;
		}
	}
	mutex_unlock(&module_mutex);
	return 0;
}

static int dim_dynamic_baseline_kernel_create(void)
{
	char digest_out[MAX_DIGEST_SIZE];
	struct dim_baseline_node *find_node = NULL;
	int ret, hash_len;
	char *kernel_rel = init_uts_ns.name.release;

	if (dim_dynamic_baseline_search(kernel_rel,
		DIM_BASELINE_KERNEL) != NULL)
		return 0;

	find_node = dim_static_baseline_search(kernel_rel,
		DIM_BASELINE_KERNEL);
	if (find_node == NULL) {
		pr_warn("can't find static baseline of kernel %s!\n",
			kernel_rel);
		dim_measure_baseline_record(kernel_rel,
			"no static baseline",
			NULL, dim_hash_get_cur_index(), 1);
	} else {
		/* output to measure log and pcr */
		dim_measure_baseline_record(kernel_rel, NULL,
			find_node->data, find_node->algo, 1);
	}

	/* create and add dynamic baseline */
	hash_len = dim_hash_get_digest_len(dim_hash_get_cur_index());
	ret = dim_calc_kernel_digest(dim_hash_get_cur_index(), digest_out);
	if (ret != hash_len) {
		pr_warn("fail to calculate kernel digest: %d\n", ret);
		return 0;
	}

	ret = dim_dynamic_baseline_create_and_add_node(kernel_rel,
		DIM_BASELINE_KERNEL, dim_hash_get_cur_index(),
		digest_out, hash_len);
	if (ret != 0) {
		pr_err("fail to add baseline of kernel: %d\n", ret);
		return 0;
	}
	return 0;
}

int dim_dynamic_baseline_build(void)
{
	int ret;

	if (dim_dynamic_baseline != NULL)
		dim_baseline_handle_destroy(dim_dynamic_baseline);

	dim_dynamic_baseline = dim_baseline_handle_create(
		dim_dynamic_baseline_compare, dim_mem_alloc, dim_mem_free);
	if (IS_ERR(dim_dynamic_baseline)) {
		ret = PTR_ERR(dim_dynamic_baseline);
		dim_dynamic_baseline = NULL;
		pr_err("fail to init dynamic baseline, reason: %d\n", ret);
		return ret;
	}

	do {
		ret = dim_dynamic_baseline_user_create();
		if (ret)
			break;

		ret = dim_dynamic_baseline_ko_create();
		if (ret)
			break;

		ret = dim_dynamic_baseline_kernel_create();
		if (ret)
			break;
	} while (0);

	if (ret) {
		pr_err("fail to create dynamic baseline: %d\n", ret);
		dim_baseline_handle_destroy(dim_dynamic_baseline);
		dim_dynamic_baseline = NULL;
		return ret;
	}

	return 0;
}

void dim_dynamic_baseline_destroy(void)
{
	dim_baseline_handle_destroy(dim_dynamic_baseline);
	dim_dynamic_baseline = NULL;
}

struct dim_baseline_node *
dim_dynamic_baseline_search(char *filename,
			    enum dim_baseline_type type)
{
	struct dim_baseline_node node = {
		.filename = filename,
		.type = type
	};

	return dim_baseline_search_node(dim_dynamic_baseline, &node);
}

#define DIM_PAGES_ARRAY_REDUNDANCY 16

struct pages_array_fun_data {
	struct dim_pages_item *array;
	unsigned long index;
};

static int filling_pgaes_array(void *start, unsigned long size, void *data)
{
	struct pages_array_fun_data *d = (struct pages_array_fun_data *)data;

	if (start == NULL || data == NULL || size == 0)
		return -EINVAL;

	d->array[d->index].pages_start = (uintptr_t)start;
	d->array[d->index].size = size;
	d->index++;

	return 0;
}

unsigned long dim_get_baseline_pages(struct dim_pages_item **pages_array)
{
	unsigned long counts = dim_mem_get_block_counts();
	struct pages_array_fun_data data;
	enum dim_core_running_status status = dim_get_running_status();

	if (counts == 0 || pages_array == NULL ||
	    status == DIM_OFF || status == DIM_NO_BASELINE ||
	    status == DIM_BASELINE_RUNNING || status == DIM_ERROR)
		return 0;

	*pages_array = vmalloc(sizeof(**pages_array) * DIM_MAX_BUDDY_NUMBER);
	if (*pages_array == NULL) {
		pr_err("alloc pages memory failed\n");
		return 0;
	}

	data.array = *pages_array;
	data.index = 0;
	dim_mem_block_traverse(filling_pgaes_array, (void *)&data);

	return data.index;
}
#if IS_MODULE(CONFIG_EULEROS_DIM_MONITOR)
EXPORT_SYMBOL(dim_get_baseline_pages);
#endif

void dim_free_pages(struct dim_pages_item **pages_array)
{
	if (pages_array == NULL || *pages_array == NULL)
		return;
	vfree(*pages_array);
	*pages_array = NULL;
}
#if IS_MODULE(CONFIG_EULEROS_DIM_MONITOR)
EXPORT_SYMBOL(dim_free_pages);
#endif
